"""
Case Management Models for AI Lawyer
SQLAlchemy models for case-specific document management
"""

from datetime import datetime
from sqlalchemy import create_engine, Column, Integer, String, DateTime, JSON, ForeignKey, Text, Boolean, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, sessionmaker
import uuid
import os

Base = declarative_base()

# Database path
DB_PATH = os.path.join(os.path.dirname(__file__), 'data', 'ai_lawyer.db')

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    email = Column(String(255), unique=True, nullable=False)
    api_key = Column(String(64), unique=True)
    subscription_tier = Column(String(50), default='free')  # free, bronze, silver, gold, platinum
    query_count = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow)
    last_login = Column(DateTime)

    # Relationships
    cases = relationship('Case', back_populates='user', cascade='all, delete-orphan')


class Case(Base):
    __tablename__ = 'cases'
    id = Column(Integer, primary_key=True)
    case_uuid = Column(String(36), unique=True, default=lambda: str(uuid.uuid4()))
    user_id = Column(Integer, ForeignKey('users.id'))

    # Case details
    name = Column(String(255), nullable=False)
    description = Column(Text)
    jurisdiction = Column(String(100))  # e.g., "CA Superior Court", "Federal - 9th Circuit"
    case_number = Column(String(100))
    area_of_law = Column(String(100))  # e.g., "personal_injury", "contract_dispute"
    status = Column(String(50), default='active')  # active, closed, archived

    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    user = relationship('User', back_populates='cases')
    documents = relationship('CaseDocument', back_populates='case', cascade='all, delete-orphan')
    queries = relationship('CaseQuery', back_populates='case', cascade='all, delete-orphan')
    metadata_record = relationship('CaseMetadata', uselist=False, back_populates='case', cascade='all, delete-orphan')


class CaseDocument(Base):
    __tablename__ = 'case_documents'
    id = Column(Integer, primary_key=True)
    case_id = Column(Integer, ForeignKey('cases.id'))

    # File info
    filename = Column(String(255), nullable=False)  # Stored filename (with hash)
    original_filename = Column(String(255), nullable=False)  # Original upload name
    file_type = Column(String(20))  # pdf, docx, txt, etc.
    file_size = Column(Integer)  # bytes
    file_hash = Column(String(64))  # MD5 hash for deduplication

    # Document classification
    document_type = Column(String(50))  # complaint, answer, motion, deposition, contract, etc.
    category = Column(String(50))  # pleadings, discovery, motions, evidence, correspondence

    # Processing status
    processing_status = Column(String(50), default='pending')  # pending, processing, completed, error
    processing_error = Column(Text)
    chunk_count = Column(Integer, default=0)

    # Vector store references
    chunk_start_index = Column(Integer)  # Starting index in case vector store
    chunk_end_index = Column(Integer)  # Ending index

    # Timestamps
    upload_date = Column(DateTime, default=datetime.utcnow)
    processed_date = Column(DateTime)

    # Relationships
    case = relationship('Case', back_populates='documents')
    chunks = relationship('DocumentChunk', back_populates='document', cascade='all, delete-orphan')


class DocumentChunk(Base):
    __tablename__ = 'document_chunks'
    id = Column(Integer, primary_key=True)
    document_id = Column(Integer, ForeignKey('case_documents.id'))

    # Chunk content
    chunk_index = Column(Integer)  # Position within document
    content = Column(Text)
    content_length = Column(Integer)

    # Page/location info
    page_number = Column(Integer)
    start_char = Column(Integer)
    end_char = Column(Integer)

    # Enrichment data
    summary = Column(Text)
    key_points = Column(JSON)
    themes = Column(JSON)
    entities = Column(JSON)  # Extracted legal entities (parties, dates, etc.)
    enriched = Column(Boolean, default=False)

    # Vector reference
    embedding_id = Column(String(64))  # Reference in vector store

    # Relationships
    document = relationship('CaseDocument', back_populates='chunks')


class CaseQuery(Base):
    __tablename__ = 'case_queries'
    id = Column(Integer, primary_key=True)
    case_id = Column(Integer, ForeignKey('cases.id'))

    # Query content
    query_text = Column(Text, nullable=False)
    query_type = Column(String(50))  # factual, legal_issue, strategy, document_search

    # Response
    response = Column(Text)
    sources = Column(JSON)  # List of document sources used
    confidence_score = Column(Float)

    # Performance
    response_time = Column(Float)  # seconds
    chunks_searched = Column(Integer)

    # Timestamp
    timestamp = Column(DateTime, default=datetime.utcnow)

    # Relationships
    case = relationship('Case', back_populates='queries')


class CaseMetadata(Base):
    __tablename__ = 'case_metadata'
    id = Column(Integer, primary_key=True)
    case_id = Column(Integer, ForeignKey('cases.id'), unique=True)

    # Extracted/entered metadata
    parties = Column(JSON)  # {"plaintiff": "John Doe", "defendant": "Acme Corp", "judge": "Hon. Smith"}
    key_dates = Column(JSON)  # {"filing_date": "2024-01-15", "trial_date": "2024-10-01"}
    claims = Column(JSON)  # List of legal claims
    defenses = Column(JSON)  # List of defenses
    damages = Column(JSON)  # Types of damages sought
    key_issues = Column(JSON)  # Legal issues identified
    key_facts = Column(JSON)  # Important facts extracted

    # Auto-generated insights
    document_summary = Column(Text)  # AI-generated case summary
    strength_assessment = Column(JSON)  # AI assessment of case strength

    # Relationships
    case = relationship('Case', back_populates='metadata_record')


# Database initialization
def init_db():
    """Initialize the database and create tables"""
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    engine = create_engine(f'sqlite:///{DB_PATH}', echo=False)
    Base.metadata.create_all(engine)
    return engine


def get_session():
    """Get a database session"""
    engine = create_engine(f'sqlite:///{DB_PATH}', echo=False)
    Session = sessionmaker(bind=engine)
    return Session()


# Initialize on import
if not os.path.exists(DB_PATH):
    init_db()