"""
Extraction Pipeline for Whisper Scribe 3.0.

Uses AI to extract entities, insights, quotes, facts, and semantic sections
from video transcriptions.
"""

import json
import uuid
import asyncio
from typing import List, Optional, Dict, Any
from datetime import datetime

from .eden_ai import chat_with_eden, calculate_chat_cost, estimate_tokens
from .project import (
    ProjectData, Entity, EntityType, Mention,
    ExtractedEntities, ExtractedInsight, InsightCategory,
    ExtractedQuote, ExtractedFact, FactCategory,
    SemanticSection, get_project, save_project
)
from .knowledge_graph import update_graph_for_project
from ..config import get_api_key, get_system_prompt


# =============================================================================
# Extraction Functions
# =============================================================================

async def extract_entities(
    transcription: str,
    api_key: str,
    provider: str = "openai"
) -> tuple[ExtractedEntities, dict]:
    """
    Extract named entities from transcription.

    Returns:
        Tuple of (ExtractedEntities, cost_info)
    """
    # Truncate if too long (keep first 12000 chars for entity extraction)
    text = transcription[:12000] if len(transcription) > 12000 else transcription

    # Get prompt from config (user-configurable)
    system_prompt = get_system_prompt("extraction_entities")
    prompt = f"""{system_prompt}

TRANSCRIPTION:
{text}"""

    messages = [{"role": "user", "content": prompt}]

    try:
        result = await chat_with_eden(
            messages=messages,
            api_key=api_key,
            providers=[provider],
            temperature=0.3,  # Lower temp for more consistent extraction
            max_tokens=3000
        )

        # Parse response
        generated = ""
        if provider in result and isinstance(result[provider], dict):
            generated = result[provider].get("generated_text", "")

        # Extract JSON from response
        entities_data = _parse_json_response(generated)

        # Convert to ExtractedEntities
        entities = ExtractedEntities(
            people=[_parse_entity(e, EntityType.PERSON) for e in entities_data.get("people", [])],
            organizations=[_parse_entity(e, EntityType.ORGANIZATION) for e in entities_data.get("organizations", [])],
            technologies=[_parse_entity(e, EntityType.TECHNOLOGY) for e in entities_data.get("technologies", [])],
            products=[_parse_entity(e, EntityType.PRODUCT) for e in entities_data.get("products", [])],
            places=[_parse_entity(e, EntityType.PLACE) for e in entities_data.get("places", [])],
            concepts=[_parse_entity(e, EntityType.CONCEPT) for e in entities_data.get("concepts", [])]
        )

        cost_info = result.get("cost_info", {})
        cost_info["operation"] = "entity_extraction"

        return entities, cost_info

    except Exception as e:
        # Return empty on error
        return ExtractedEntities(), {"error": str(e)}


async def extract_insights(
    transcription: str,
    api_key: str,
    provider: str = "openai"
) -> tuple[List[ExtractedInsight], dict]:
    """
    Extract key insights from transcription.

    Returns:
        Tuple of (list of insights, cost_info)
    """
    text = transcription[:12000] if len(transcription) > 12000 else transcription

    # Get prompt from config (user-configurable)
    system_prompt = get_system_prompt("extraction_insights")
    prompt = f"""{system_prompt}

TRANSCRIPTION:
{text}"""

    messages = [{"role": "user", "content": prompt}]

    try:
        result = await chat_with_eden(
            messages=messages,
            api_key=api_key,
            providers=[provider],
            temperature=0.5,
            max_tokens=2000
        )

        generated = ""
        if provider in result and isinstance(result[provider], dict):
            generated = result[provider].get("generated_text", "")

        insights_data = _parse_json_response(generated)

        insights = []
        for item in insights_data.get("insights", []):
            category = _parse_insight_category(item.get("category", "takeaway"))
            insights.append(ExtractedInsight(
                id=str(uuid.uuid4()),
                content=item.get("content", ""),
                category=category,
                confidence=0.9
            ))

        cost_info = result.get("cost_info", {})
        cost_info["operation"] = "insight_extraction"

        return insights, cost_info

    except Exception as e:
        return [], {"error": str(e)}


async def extract_quotes(
    transcription: str,
    api_key: str,
    provider: str = "openai"
) -> tuple[List[ExtractedQuote], dict]:
    """
    Extract notable quotes from transcription.

    Returns:
        Tuple of (list of quotes, cost_info)
    """
    text = transcription[:12000] if len(transcription) > 12000 else transcription

    # Get prompt from config (user-configurable)
    system_prompt = get_system_prompt("extraction_quotes")
    prompt = f"""{system_prompt}

TRANSCRIPTION:
{text}"""

    messages = [{"role": "user", "content": prompt}]

    try:
        result = await chat_with_eden(
            messages=messages,
            api_key=api_key,
            providers=[provider],
            temperature=0.4,
            max_tokens=2000
        )

        generated = ""
        if provider in result and isinstance(result[provider], dict):
            generated = result[provider].get("generated_text", "")

        quotes_data = _parse_json_response(generated)

        quotes = []
        for item in quotes_data.get("quotes", []):
            quotes.append(ExtractedQuote(
                id=str(uuid.uuid4()),
                text=item.get("text", ""),
                speaker=item.get("speaker"),
                context=item.get("context", "")
            ))

        cost_info = result.get("cost_info", {})
        cost_info["operation"] = "quote_extraction"

        return quotes, cost_info

    except Exception as e:
        return [], {"error": str(e)}


async def extract_facts(
    transcription: str,
    api_key: str,
    provider: str = "openai"
) -> tuple[List[ExtractedFact], dict]:
    """
    Extract facts and statistics from transcription.

    Returns:
        Tuple of (list of facts, cost_info)
    """
    text = transcription[:12000] if len(transcription) > 12000 else transcription

    # Get prompt from config (user-configurable)
    system_prompt = get_system_prompt("extraction_facts")
    prompt = f"""{system_prompt}

TRANSCRIPTION:
{text}"""

    messages = [{"role": "user", "content": prompt}]

    try:
        result = await chat_with_eden(
            messages=messages,
            api_key=api_key,
            providers=[provider],
            temperature=0.3,
            max_tokens=2000
        )

        generated = ""
        if provider in result and isinstance(result[provider], dict):
            generated = result[provider].get("generated_text", "")

        facts_data = _parse_json_response(generated)

        facts = []
        for item in facts_data.get("facts", []):
            category = _parse_fact_category(item.get("category", "claim"))
            facts.append(ExtractedFact(
                id=str(uuid.uuid4()),
                statement=item.get("statement", ""),
                category=category,
                source=item.get("source"),
                confidence=0.85
            ))

        cost_info = result.get("cost_info", {})
        cost_info["operation"] = "fact_extraction"

        return facts, cost_info

    except Exception as e:
        return [], {"error": str(e)}


async def extract_sections(
    transcription: str,
    api_key: str,
    provider: str = "openai",
    duration_seconds: int = 0
) -> tuple[List[SemanticSection], dict]:
    """
    Extract semantic sections from transcription.

    Returns:
        Tuple of (list of sections, cost_info)
    """
    text = transcription[:15000] if len(transcription) > 15000 else transcription

    # Get prompt from config (user-configurable)
    system_prompt = get_system_prompt("extraction_sections")
    prompt = f"""{system_prompt}

TRANSCRIPTION:
{text}"""

    messages = [{"role": "user", "content": prompt}]

    try:
        result = await chat_with_eden(
            messages=messages,
            api_key=api_key,
            providers=[provider],
            temperature=0.4,
            max_tokens=2000
        )

        generated = ""
        if provider in result and isinstance(result[provider], dict):
            generated = result[provider].get("generated_text", "")

        sections_data = _parse_json_response(generated)

        sections = []
        raw_sections = sections_data.get("sections", [])
        num_sections = len(raw_sections)

        for i, item in enumerate(raw_sections):
            # Estimate timestamps based on position
            position = item.get("approximate_position", "middle")
            start_time, end_time = _estimate_timestamps(
                position, i, num_sections, duration_seconds
            )

            sections.append(SemanticSection(
                id=str(uuid.uuid4()),
                title=item.get("title", f"Section {i+1}"),
                content="",  # Content can be filled from transcription later
                start_time=start_time,
                end_time=end_time,
                topics=item.get("topics", []),
                summary=item.get("summary", "")
            ))

        cost_info = result.get("cost_info", {})
        cost_info["operation"] = "section_extraction"

        return sections, cost_info

    except Exception as e:
        return [], {"error": str(e)}


# =============================================================================
# Extraction Orchestrator
# =============================================================================

async def run_full_extraction(
    video_id: str,
    transcription: str,
    api_key: str,
    provider: str = "openai",
    duration_seconds: int = 0,
    extract_entities_flag: bool = True,
    extract_insights_flag: bool = True,
    extract_quotes_flag: bool = True,
    extract_facts_flag: bool = True,
    extract_sections_flag: bool = True
) -> Dict[str, Any]:
    """
    Run the full extraction pipeline on a transcription.

    Args:
        video_id: The video/project ID
        transcription: The full transcription text
        api_key: Eden.ai API key
        provider: AI provider to use
        duration_seconds: Video duration for timestamp estimation
        extract_*_flag: Flags to control which extractions to run

    Returns:
        Dict with extraction results and costs
    """
    results = {
        "video_id": video_id,
        "status": "in_progress",
        "entities": None,
        "insights": [],
        "quotes": [],
        "facts": [],
        "sections": [],
        "costs": [],
        "total_cost": 0.0,
        "errors": []
    }

    # Run extractions in parallel where possible
    tasks = []

    if extract_entities_flag:
        tasks.append(("entities", extract_entities(transcription, api_key, provider)))

    if extract_insights_flag:
        tasks.append(("insights", extract_insights(transcription, api_key, provider)))

    if extract_quotes_flag:
        tasks.append(("quotes", extract_quotes(transcription, api_key, provider)))

    if extract_facts_flag:
        tasks.append(("facts", extract_facts(transcription, api_key, provider)))

    if extract_sections_flag:
        tasks.append(("sections", extract_sections(transcription, api_key, provider, duration_seconds)))

    # Execute all tasks
    if tasks:
        task_results = await asyncio.gather(
            *[task[1] for task in tasks],
            return_exceptions=True
        )

        # Process results
        total_cost = 0.0
        for i, (name, _) in enumerate(tasks):
            task_result = task_results[i]

            if isinstance(task_result, Exception):
                results["errors"].append({
                    "extraction": name,
                    "error": str(task_result)
                })
            else:
                data, cost_info = task_result
                results[name] = data if name == "entities" else data

                if cost_info and "cost_usd" in cost_info:
                    results["costs"].append(cost_info)
                    total_cost += cost_info.get("cost_usd", 0)
                elif cost_info and "error" in cost_info:
                    results["errors"].append({
                        "extraction": name,
                        "error": cost_info["error"]
                    })

        results["total_cost"] = round(total_cost, 6)

    results["status"] = "complete" if not results["errors"] else "partial"

    return results


async def run_extraction_and_save(
    video_id: str,
    api_key: str = None,
    provider: str = "openai",
    **kwargs
) -> Dict[str, Any]:
    """
    Run extraction pipeline and save results to project.

    Args:
        video_id: The project/video ID
        api_key: Eden.ai API key (uses config if not provided)
        provider: AI provider to use
        **kwargs: Additional flags for run_full_extraction

    Returns:
        Extraction results dict
    """
    # Get project
    project = get_project(video_id)
    if not project:
        return {"error": "Project not found", "video_id": video_id}

    # Get API key
    if not api_key:
        api_key = get_api_key("eden")
        if not api_key:
            return {"error": "No Eden.ai API key configured", "video_id": video_id}

    # Update status
    project.extraction_status = "in_progress"
    save_project(project)

    try:
        # Run extraction
        results = await run_full_extraction(
            video_id=video_id,
            transcription=project.transcription,
            api_key=api_key,
            provider=provider,
            duration_seconds=project.duration,
            **kwargs
        )

        # Update project with results
        if results.get("entities"):
            project.entities = results["entities"]

        if results.get("insights"):
            project.insights = results["insights"]

        if results.get("quotes"):
            project.quotes = results["quotes"]

        if results.get("facts"):
            project.facts = results["facts"]

        if results.get("sections"):
            project.sections = results["sections"]

        # Update costs
        for cost in results.get("costs", []):
            project.costs["entries"].append(cost)
            project.costs["total_cost"] += cost.get("cost_usd", 0)
        project.costs["total_cost"] = round(project.costs["total_cost"], 6)

        # Extract topics from sections and entities
        topics = set(project.shared_topics)
        for section in project.sections:
            topics.update(section.topics)
        for entity in project.entities.technologies:
            topics.add(entity.name)
        for entity in project.entities.concepts:
            topics.add(entity.name)
        project.shared_topics = list(topics)

        # Update status
        project.extraction_status = "complete" if results["status"] == "complete" else "partial"
        project.processed_at = datetime.now().isoformat()

        # Save project
        save_project(project)

        # Update knowledge graph
        update_graph_for_project(project)

        return results

    except Exception as e:
        project.extraction_status = "error"
        save_project(project)
        return {"error": str(e), "video_id": video_id}


# =============================================================================
# Helper Functions
# =============================================================================

def _parse_json_response(text: str) -> dict:
    """Parse JSON from AI response, handling common issues."""
    if not text:
        return {}

    # Try to find JSON in the response
    text = text.strip()

    # Remove markdown code blocks if present
    if text.startswith("```json"):
        text = text[7:]
    elif text.startswith("```"):
        text = text[3:]
    if text.endswith("```"):
        text = text[:-3]

    text = text.strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        # Try to find JSON object in text
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            try:
                return json.loads(text[start:end])
            except json.JSONDecodeError:
                pass
        return {}


def _parse_entity(data: dict, entity_type: EntityType) -> Entity:
    """Parse entity data into Entity object."""
    mentions = []
    for mention_text in data.get("mentions", []):
        if isinstance(mention_text, str):
            mentions.append(Mention(timestamp=0.0, context=mention_text))
        elif isinstance(mention_text, dict):
            mentions.append(Mention(
                timestamp=mention_text.get("timestamp", 0.0),
                context=mention_text.get("context", "")
            ))

    return Entity(
        id=str(uuid.uuid4()),
        name=data.get("name", "Unknown"),
        type=entity_type,
        description=data.get("description", ""),
        mentions=mentions,
        wikipedia_url=data.get("wikipedia_url")
    )


def _parse_insight_category(category_str: str) -> InsightCategory:
    """Parse insight category string to enum."""
    category_map = {
        "takeaway": InsightCategory.TAKEAWAY,
        "aha": InsightCategory.AHA,
        "wisdom": InsightCategory.WISDOM,
        "action": InsightCategory.ACTION,
        "warning": InsightCategory.WARNING,
        "definition": InsightCategory.DEFINITION
    }
    return category_map.get(category_str.lower(), InsightCategory.TAKEAWAY)


def _parse_fact_category(category_str: str) -> FactCategory:
    """Parse fact category string to enum."""
    category_map = {
        "statistic": FactCategory.STATISTIC,
        "definition": FactCategory.DEFINITION,
        "claim": FactCategory.CLAIM,
        "date": FactCategory.DATE,
        "comparison": FactCategory.COMPARISON
    }
    return category_map.get(category_str.lower(), FactCategory.CLAIM)


def _estimate_timestamps(
    position: str,
    index: int,
    total_sections: int,
    duration_seconds: int
) -> tuple[float, float]:
    """Estimate start and end timestamps for a section."""
    if duration_seconds <= 0:
        return 0.0, 0.0

    # Use position hint and index to estimate
    position_map = {
        "beginning": 0.0,
        "early": 0.2,
        "middle": 0.5,
        "late": 0.7,
        "end": 0.9
    }

    # Calculate based on index if we have sections
    if total_sections > 1:
        section_duration = duration_seconds / total_sections
        start_time = index * section_duration
        end_time = (index + 1) * section_duration
    else:
        # Use position hint
        pos_fraction = position_map.get(position.lower(), 0.5)
        section_duration = duration_seconds * 0.2  # Assume ~20% of video
        start_time = pos_fraction * duration_seconds
        end_time = min(start_time + section_duration, duration_seconds)

    return round(start_time, 1), round(end_time, 1)


# =============================================================================
# Sync Wrappers
# =============================================================================

def run_extraction_sync(video_id: str, **kwargs) -> Dict[str, Any]:
    """Synchronous wrapper for run_extraction_and_save."""
    import asyncio

    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # We're in an async context, create a new task
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(
                    asyncio.run,
                    run_extraction_and_save(video_id, **kwargs)
                )
                return future.result()
        else:
            return loop.run_until_complete(run_extraction_and_save(video_id, **kwargs))
    except RuntimeError:
        # No event loop, create one
        return asyncio.run(run_extraction_and_save(video_id, **kwargs))
