"""
Annotation Manager - Manages trade annotations and test case generation

Handles storage, retrieval, and test case generation from manual trade annotations.
"""

import json
import uuid
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, asdict
import logging

logger = logging.getLogger(__name__)


@dataclass
class TradeAnnotation:
    """Represents a manually marked trade"""
    annotation_id: str
    symbol: str
    timeframe: str
    entry: Dict[str, Any]  # {timestamp, price, index}
    exit: Dict[str, Any]   # {timestamp, price, index}
    direction: str  # 'LONG' or 'SHORT'
    profit_loss_pct: float
    notes: str = ""
    created_at: str = None
    market_context: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now().isoformat()
        if self.market_context is None:
            self.market_context = {}


class AnnotationManager:
    """Manages trade annotations and test case generation"""
    
    def __init__(self, storage_path: str = "ANNOTATE/data/annotations"):
        """Initialize annotation manager"""
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)
        
        self.annotations_file = self.storage_path / "annotations_db.json"
        self.test_cases_dir = self.storage_path.parent / "test_cases"
        self.test_cases_dir.mkdir(parents=True, exist_ok=True)
        
        self.annotations_db = self._load_annotations()
        
        logger.info(f"AnnotationManager initialized with storage: {self.storage_path}")
    
    def _load_annotations(self) -> Dict[str, List[Dict]]:
        """Load annotations from storage"""
        if self.annotations_file.exists():
            try:
                with open(self.annotations_file, 'r') as f:
                    data = json.load(f)
                    logger.info(f"Loaded {len(data.get('annotations', []))} annotations")
                    return data
            except Exception as e:
                logger.error(f"Error loading annotations: {e}")
                return {"annotations": [], "metadata": {}}
        else:
            return {"annotations": [], "metadata": {}}
    
    def _save_annotations(self):
        """Save annotations to storage"""
        try:
            # Update metadata
            self.annotations_db["metadata"] = {
                "total_annotations": len(self.annotations_db["annotations"]),
                "last_updated": datetime.now().isoformat()
            }
            
            with open(self.annotations_file, 'w') as f:
                json.dump(self.annotations_db, f, indent=2)
            
            logger.info(f"Saved {len(self.annotations_db['annotations'])} annotations")
        except Exception as e:
            logger.error(f"Error saving annotations: {e}")
            raise
    
    def create_annotation(self, entry_point: Dict, exit_point: Dict, 
                         symbol: str, timeframe: str, 
                         entry_market_state: Dict = None,
                         exit_market_state: Dict = None) -> TradeAnnotation:
        """Create new trade annotation"""
        # Calculate direction and P&L
        entry_price = entry_point['price']
        exit_price = exit_point['price']
        
        if exit_price > entry_price:
            direction = 'LONG'
            profit_loss_pct = ((exit_price - entry_price) / entry_price) * 100
        else:
            direction = 'SHORT'
            profit_loss_pct = ((entry_price - exit_price) / entry_price) * 100
        
        annotation = TradeAnnotation(
            annotation_id=str(uuid.uuid4()),
            symbol=symbol,
            timeframe=timeframe,
            entry=entry_point,
            exit=exit_point,
            direction=direction,
            profit_loss_pct=profit_loss_pct
        )
        
        logger.info(f"Created annotation: {annotation.annotation_id} ({direction}, {profit_loss_pct:.2f}%)")
        return annotation
    
    def save_annotation(self, annotation: TradeAnnotation):
        """Save annotation to storage"""
        # Convert to dict
        ann_dict = asdict(annotation)
        
        # Add to database
        self.annotations_db["annotations"].append(ann_dict)
        
        # Save to file
        self._save_annotations()
        
        logger.info(f"Saved annotation: {annotation.annotation_id}")
    
    def get_annotations(self, symbol: str = None, 
                       timeframe: str = None) -> List[TradeAnnotation]:
        """Retrieve annotations with optional filtering"""
        annotations = self.annotations_db.get("annotations", [])
        
        # Filter by symbol
        if symbol:
            annotations = [a for a in annotations if a.get('symbol') == symbol]
        
        # Filter by timeframe
        if timeframe:
            annotations = [a for a in annotations if a.get('timeframe') == timeframe]
        
        # Convert to TradeAnnotation objects
        result = []
        for ann_dict in annotations:
            try:
                annotation = TradeAnnotation(**ann_dict)
                result.append(annotation)
            except Exception as e:
                logger.error(f"Error converting annotation: {e}")
        
        return result
    
    def delete_annotation(self, annotation_id: str):
        """Delete annotation"""
        original_count = len(self.annotations_db["annotations"])
        self.annotations_db["annotations"] = [
            a for a in self.annotations_db["annotations"] 
            if a.get('annotation_id') != annotation_id
        ]
        
        if len(self.annotations_db["annotations"]) < original_count:
            self._save_annotations()
            logger.info(f"Deleted annotation: {annotation_id}")
        else:
            logger.warning(f"Annotation not found: {annotation_id}")
    
    def generate_test_case(self, annotation: TradeAnnotation, data_provider=None) -> Dict:
        """
        Generate test case from annotation in realtime format
        
        Args:
            annotation: TradeAnnotation object
            data_provider: Optional DataProvider instance to fetch market context
            
        Returns:
            Test case dictionary in realtime format
        """
        test_case = {
            "test_case_id": f"annotation_{annotation.annotation_id}",
            "symbol": annotation.symbol,
            "timestamp": annotation.entry['timestamp'],
            "action": "BUY" if annotation.direction == "LONG" else "SELL",
            "market_state": {},
            "expected_outcome": {
                "direction": annotation.direction,
                "profit_loss_pct": annotation.profit_loss_pct,
                "holding_period_seconds": self._calculate_holding_period(annotation),
                "exit_price": annotation.exit['price'],
                "entry_price": annotation.entry['price']
            },
            "annotation_metadata": {
                "annotator": "manual",
                "confidence": 1.0,
                "notes": annotation.notes,
                "created_at": annotation.created_at,
                "timeframe": annotation.timeframe
            }
        }
        
        # Populate market state if data_provider is available
        if data_provider and annotation.market_context:
            test_case["market_state"] = annotation.market_context
        elif data_provider:
            # Fetch market state at entry time
            try:
                entry_time = datetime.fromisoformat(annotation.entry['timestamp'].replace('Z', '+00:00'))
                
                # Fetch OHLCV data for all timeframes
                timeframes = ['1s', '1m', '1h', '1d']
                market_state = {}
                
                for tf in timeframes:
                    df = data_provider.get_historical_data(
                        symbol=annotation.symbol,
                        timeframe=tf,
                        limit=100
                    )
                    
                    if df is not None and not df.empty:
                        # Filter to data before entry time
                        df = df[df.index <= entry_time]
                        
                        if not df.empty:
                            # Convert to list format
                            market_state[f'ohlcv_{tf}'] = {
                                'timestamps': df.index.strftime('%Y-%m-%d %H:%M:%S').tolist(),
                                'open': df['open'].tolist(),
                                'high': df['high'].tolist(),
                                'low': df['low'].tolist(),
                                'close': df['close'].tolist(),
                                'volume': df['volume'].tolist()
                            }
                
                test_case["market_state"] = market_state
                logger.info(f"Populated market state with {len(market_state)} timeframes")
                
            except Exception as e:
                logger.error(f"Error fetching market state: {e}")
                test_case["market_state"] = {}
        else:
            test_case["market_state"] = {}
        
        # Save test case to file
        test_case_file = self.test_cases_dir / f"{test_case['test_case_id']}.json"
        with open(test_case_file, 'w') as f:
            json.dump(test_case, f, indent=2)
        
        logger.info(f"Generated test case: {test_case['test_case_id']}")
        return test_case
    
    def _calculate_holding_period(self, annotation: TradeAnnotation) -> float:
        """Calculate holding period in seconds"""
        try:
            entry_time = datetime.fromisoformat(annotation.entry['timestamp'].replace('Z', '+00:00'))
            exit_time = datetime.fromisoformat(annotation.exit['timestamp'].replace('Z', '+00:00'))
            return (exit_time - entry_time).total_seconds()
        except Exception as e:
            logger.error(f"Error calculating holding period: {e}")
            return 0.0
    
    def export_annotations(self, annotations: List[TradeAnnotation] = None, 
                          format_type: str = 'json') -> Path:
        """Export annotations to file"""
        if annotations is None:
            annotations = self.get_annotations()
        
        # Convert to dicts
        export_data = [asdict(ann) for ann in annotations]
        
        # Create export file
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        export_file = self.storage_path / f"export_{timestamp}.{format_type}"
        
        if format_type == 'json':
            with open(export_file, 'w') as f:
                json.dump(export_data, f, indent=2)
        elif format_type == 'csv':
            import csv
            with open(export_file, 'w', newline='') as f:
                if export_data:
                    writer = csv.DictWriter(f, fieldnames=export_data[0].keys())
                    writer.writeheader()
                    writer.writerows(export_data)
        
        logger.info(f"Exported {len(annotations)} annotations to {export_file}")
        return export_file