"""Persistence layer for LLM requests and responses."""

from datetime import datetime
from typing import Optional, List
from uuid import UUID

from sqlalchemy import create_engine, Column, String, Integer, Float, Boolean, DateTime, Text, JSON
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session

from .models import LLMRequest, LLMResponse, TokenUsage, ProvenanceMetadata

Base = declarative_base()


class LLMRequestRecord(Base):
    """Database model for LLM requests."""
    
    __tablename__ = "llm_requests"
    
    id = Column(PG_UUID(as_uuid=True), primary_key=True)
    content_hash = Column(String(64), index=True, nullable=False)
    
    model = Column(String(100), nullable=False)
    messages = Column(JSON, nullable=False)
    temperature = Column(Float, nullable=False)
    max_tokens = Column(Integer, nullable=False)
    top_p = Column(Float, nullable=True)
    stop_sequences = Column(JSON, nullable=True)
    
    # Provenance
    workflow_id = Column(String(100), index=True, nullable=True)
    chunk_id = Column(String(100), index=True, nullable=True)
    section_id = Column(String(100), nullable=True)
    source_id = Column(String(100), nullable=True)
    character_range = Column(JSON, nullable=True)
    caller_id = Column(String(100), nullable=True)
    metadata_tags = Column(JSON, nullable=True)
    
    extra_params = Column(JSON, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)


class LLMResponseRecord(Base):
    """Database model for LLM responses."""
    
    __tablename__ = "llm_responses"
    
    id = Column(PG_UUID(as_uuid=True), primary_key=True)
    request_id = Column(PG_UUID(as_uuid=True), index=True, nullable=False)
    
    content = Column(Text, nullable=False)
    
    # Token usage
    prompt_tokens = Column(Integer, nullable=True)
    completion_tokens = Column(Integer, nullable=True)
    total_tokens = Column(Integer, nullable=True)
    
    # Cache info
    from_cache = Column(Boolean, default=False, nullable=False)
    cache_key = Column(String(100), nullable=True)
    
    # Provider metadata
    provider = Column(String(50), nullable=False)
    model = Column(String(100), nullable=False)
    provider_request_id = Column(String(200), nullable=True)
    latency_ms = Column(Float, nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)


class LLMPersistence:
    """Persistence manager for LLM requests and responses."""
    
    def __init__(self, database_url: str = "sqlite:///llm_gateway.db"):
        self.engine = create_engine(database_url, pool_pre_ping=True)
        Base.metadata.create_all(self.engine)
        self.SessionLocal = sessionmaker(bind=self.engine)
    
    def _get_session(self) -> Session:
        """Get database session."""
        return self.SessionLocal()
    
    def save_request(self, request: LLMRequest, content_hash: str) -> None:
        """Save LLM request to database."""
        with self._get_session() as session:
            record = LLMRequestRecord(
                id=request.id,
                content_hash=content_hash,
                model=request.model,
                messages=request.messages,
                temperature=request.temperature,
                max_tokens=request.max_tokens,
                top_p=request.top_p,
                stop_sequences=request.stop_sequences,
                workflow_id=request.metadata.workflow_id,
                chunk_id=request.metadata.chunk_id,
                section_id=request.metadata.section_id,
                source_id=request.metadata.source_id,
                character_range=list(request.metadata.character_range) if request.metadata.character_range else None,
                caller_id=request.metadata.caller_id,
                metadata_tags=request.metadata.tags,
                extra_params=request.extra_params,
                created_at=request.created_at
            )
            session.add(record)
            session.commit()
    
    def save_response(self, response: LLMResponse) -> None:
        """Save LLM response to database."""
        with self._get_session() as session:
            record = LLMResponseRecord(
                id=response.id,
                request_id=response.request_id,
                content=response.content,
                prompt_tokens=response.usage.prompt_tokens if response.usage else None,
                completion_tokens=response.usage.completion_tokens if response.usage else None,
                total_tokens=response.usage.total_tokens if response.usage else None,
                from_cache=response.from_cache,
                cache_key=response.cache_key,
                provider=response.provider,
                model=response.model,
                provider_request_id=response.provider_request_id,
                latency_ms=response.latency_ms,
                created_at=response.created_at
            )
            session.add(record)
            session.commit()
    
    def get_request(self, request_id: UUID) -> Optional[LLMRequest]:
        """Get request by ID."""
        with self._get_session() as session:
            record = session.query(LLMRequestRecord).filter_by(id=request_id).first()
            if not record:
                return None
            
            return LLMRequest(
                id=record.id,
                model=record.model,
                messages=record.messages,
                temperature=record.temperature,
                max_tokens=record.max_tokens,
                top_p=record.top_p,
                stop_sequences=record.stop_sequences,
                metadata=ProvenanceMetadata(
                    workflow_id=record.workflow_id,
                    chunk_id=record.chunk_id,
                    section_id=record.section_id,
                    source_id=record.source_id,
                    character_range=tuple(record.character_range) if record.character_range else None,
                    caller_id=record.caller_id,
                    tags=record.metadata_tags or {}
                ),
                extra_params=record.extra_params or {},
                created_at=record.created_at
            )
    
    def get_response(self, response_id: UUID) -> Optional[LLMResponse]:
        """Get response by ID."""
        with self._get_session() as session:
            record = session.query(LLMResponseRecord).filter_by(id=response_id).first()
            if not record:
                return None
            
            # Get associated request for metadata
            request_record = session.query(LLMRequestRecord).filter_by(id=record.request_id).first()
            metadata = ProvenanceMetadata()
            if request_record:
                metadata = ProvenanceMetadata(
                    workflow_id=request_record.workflow_id,
                    chunk_id=request_record.chunk_id,
                    section_id=request_record.section_id,
                    source_id=request_record.source_id,
                    character_range=tuple(request_record.character_range) if request_record.character_range else None,
                    caller_id=request_record.caller_id,
                    tags=request_record.metadata_tags or {}
                )
            
            usage = None
            if record.total_tokens:
                usage = TokenUsage(
                    prompt_tokens=record.prompt_tokens or 0,
                    completion_tokens=record.completion_tokens or 0,
                    total_tokens=record.total_tokens
                )
            
            return LLMResponse(
                id=record.id,
                request_id=record.request_id,
                content=record.content,
                usage=usage,
                from_cache=record.from_cache,
                cache_key=record.cache_key,
                provider=record.provider,
                model=record.model,
                provider_request_id=record.provider_request_id,
                latency_ms=record.latency_ms,
                metadata=metadata,
                created_at=record.created_at
            )
    
    def get_requests_by_metadata(
        self,
        workflow_id: Optional[str] = None,
        chunk_id: Optional[str] = None,
        limit: int = 100
    ) -> List[LLMRequest]:
        """Query requests by metadata."""
        with self._get_session() as session:
            query = session.query(LLMRequestRecord)
            
            if workflow_id:
                query = query.filter_by(workflow_id=workflow_id)
            if chunk_id:
                query = query.filter_by(chunk_id=chunk_id)
            
            records = query.order_by(LLMRequestRecord.created_at.desc()).limit(limit).all()
            
            return [
                LLMRequest(
                    id=r.id,
                    model=r.model,
                    messages=r.messages,
                    temperature=r.temperature,
                    max_tokens=r.max_tokens,
                    metadata=ProvenanceMetadata(
                        workflow_id=r.workflow_id,
                        chunk_id=r.chunk_id,
                        section_id=r.section_id,
                        source_id=r.source_id,
                        tags=r.metadata_tags or {}
                    ),
                    created_at=r.created_at
                )
                for r in records
            ]
    
    def get_token_usage_summary(
        self,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None
    ) -> dict:
        """Get aggregated token usage statistics."""
        with self._get_session() as session:
            query = session.query(
                LLMResponseRecord.model,
                LLMResponseRecord.provider
            )
            
            if start_date:
                query = query.filter(LLMResponseRecord.created_at >= start_date)
            if end_date:
                query = query.filter(LLMResponseRecord.created_at <= end_date)
            
            records = query.all()
            
            summary = {
                "total_requests": len(records),
                "by_model": {},
                "by_provider": {}
            }
            
            return summary
