"""
DomainAnalyzer — Extracts a world model from a problem statement via LLM.

Port of src/afs/worldmodel/DomainAnalyzer.js to Python.

Pipeline (parallelized):
  Phase 1: Concepts + domain + objectives + tradeoffs (fast, gpt-4.1-nano)
  Phase 2 (parallel):
    A: Relationships + truth tables
    B: Attractors + reference points

Output is structured data ready for WorldModelTranslator consumption.
"""

import asyncio
import json
import time
from typing import Optional, Callable, Any

try:
    from openai import AsyncOpenAI
except ImportError:
    AsyncOpenAI = None


class DomainAnalyzer:
    def __init__(self, **options):
        if AsyncOpenAI is None:
            raise ImportError("openai package required: pip install openai")
        self.client = options.get("client") or AsyncOpenAI()
        self.model = options.get("model", "gpt-4.1-mini")
        self.max_influences = options.get("max_influences", 15)
        self.min_influences = options.get("min_influences", 4)

    async def analyze(
        self,
        problem_statement: str,
        options: Optional[dict] = None,
        on_progress: Optional[Callable] = None,
    ) -> dict:
        """
        Analyze a problem domain and produce a structured world model.

        Returns normalized world model dict with concepts, relationships,
        attractors, referencePoints, and confidence assessment.
        """
        options = options or {}
        max_inf = options.get("max_influences", self.max_influences)
        user_context = self._build_user_context(problem_statement, options)
        report = on_progress or (lambda *a: None)

        t0 = time.time()

        # Phase 1: Concepts (fast model)
        phase1_model = "gpt-4.1-nano"
        phase1 = await self._run_phase1(user_context, max_inf, phase1_model)
        t1 = time.time()
        print(f"[DomainAnalyzer] Phase 1 ({phase1_model}): {(t1-t0)*1000:.0f}ms — {len(phase1.get('concepts', []))} concepts")
        report("phase1_done", f"{len(phase1.get('concepts', []))} concepts extracted in {t1-t0:.1f}s")

        # Phase 2: Parallel — relationships AND attractors
        concept_summary = self._build_concept_summary(phase1)
        report("phase2_start", f"Building relationships & attractors for {len(phase1.get('concepts', []))} concepts (parallel)...")

        phase2_rels, phase2_attr = await asyncio.gather(
            self._run_phase2_relationships(user_context, concept_summary, phase1),
            self._run_phase2_attractors(user_context, concept_summary, phase1),
        )
        t2 = time.time()
        rels = phase2_rels.get("relationships", [])
        attrs = phase2_attr.get("attractors", [])
        print(f"[DomainAnalyzer] Phase 2 ({self.model}, parallel): {(t2-t1)*1000:.0f}ms — {len(rels)} relationships, {len(attrs)} attractors")
        print(f"[DomainAnalyzer] Total: {(t2-t0)*1000:.0f}ms")
        report("phase2_done", f"{len(rels)} relationships, {len(attrs)} attractors in {t2-t1:.1f}s")

        # Reconcile concept names (fuzzy match Phase 2 → Phase 1)
        valid_names = {c["name"] for c in phase1.get("concepts", [])}
        valid_names_list = list(valid_names)

        def match_concept_name(name: str) -> Optional[str]:
            if name in valid_names:
                return name
            lower = name.lower()
            for vn in valid_names_list:
                if vn.lower() == lower:
                    return vn
            for vn in valid_names_list:
                if vn in name or name in vn:
                    return vn
            name_words = [w for w in name.replace("_", " ").split() if len(w) > 2]
            for vn in valid_names_list:
                vn_words = [w for w in vn.replace("_", " ").split() if len(w) > 2]
                overlap = [w for w in name_words if any(vw in w or w in vw for vw in vn_words)]
                if overlap and len(overlap) >= min(len(name_words), len(vn_words)) * 0.5:
                    return vn
            return None

        reconciled_rels = []
        for r in rels:
            from_match = match_concept_name(r.get("from", ""))
            to_match = match_concept_name(r.get("to", ""))
            if not from_match or not to_match:
                continue
            r["from"] = from_match
            r["to"] = to_match
            reconciled_rels.append(r)

        reconciled_attractors = []
        for a in phase2_attr.get("attractors", []):
            pos = a.get("position", {})
            cleaned = {}
            for k, v in pos.items():
                match = match_concept_name(k)
                if match:
                    cleaned[match] = v
            reconciled_attractors.append({**a, "position": cleaned})

        reconciled_ref_points = []
        for rp in phase2_attr.get("referencePoints", []):
            pos = rp.get("position", {})
            cleaned = {}
            for k, v in pos.items():
                match = match_concept_name(k)
                if match:
                    cleaned[match] = v
            reconciled_ref_points.append({**rp, "position": cleaned})

        merged = {
            "domain": phase1.get("domain", "general"),
            "problemRestatement": phase1.get("problemRestatement", problem_statement),
            "objectives": phase1.get("objectives", []),
            "concepts": phase1.get("concepts", []),
            "relationships": reconciled_rels,
            "tradeoffs": phase1.get("tradeoffs", []),
            "attractors": reconciled_attractors,
            "referencePoints": reconciled_ref_points,
            "confidenceAssessment": phase2_attr.get("confidenceAssessment", {
                "overall": "medium", "weakestLink": "unknown", "dataGaps": []
            }),
        }

        return self._validate_and_normalize(merged, problem_statement)

    # ─── Phase 1: Concepts ─────────────────────────────────────

    async def _run_phase1(self, user_context: str, max_influences: int, model: str) -> dict:
        system_prompt = f"""You are a domain expert building a quantitative simulation model. Identify the real-world forces, constraints, resources, and dynamics that govern this system.

Produce a JSON response with this structure:

{{
    "domain": "<one word domain name>",
    "problemRestatement": "<clear 1-3 sentence restatement>",
    "objectives": ["<primary objective>", "<secondary objective>"],
    "concepts": [
        {{
            "name": "<snake_case identifier>",
            "displayName": "<Human Readable Name>",
            "description": "<2-3 sentences: what this represents in reality, what it controls, what it responds to>",
            "nature": "<resource|constraint|force|behavior|outcome>",
            "category": "<domain-specific grouping>",
            "controllability": "<direct|indirect|uncontrollable>",
            "unitOfMeasure": "<real-world unit>",
            "validRange": [<min>, <max>],
            "displayRange": [<min>, <max>],
            "displayUnit": "<display unit string>",
            "historicalBaseline": <typical/median value>,
            "dynamics": {{
                "inertia": "<low|medium|high>",
                "timescale": "<days|weeks|months|quarters|years>",
                "volatility": <0.0-1.0>
            }}
        }}
    ],
    "tradeoffs": [
        {{
            "concepts": ["<concept_a>", "<concept_b>"],
            "description": "<why these are fundamentally in tension>",
            "severity": "<mild|moderate|severe>"
        }}
    ]
}}

RULES:
1. CONCEPTS must be real-world forces with measurable units and ranges — not abstract labels.
2. Generate {self.min_influences} to {max_influences} concepts. Prefer fewer, more meaningful concepts.
3. Every problem has TRADE-OFFS. Identify at least 2 pairs of concepts in fundamental tension.
4. Focus on QUALITY of concepts — detailed descriptions, accurate ranges, real units."""

        response = await self.client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_context},
            ],
            response_format={"type": "json_object"},
            temperature=0.5,
            max_tokens=6000,
        )
        return json.loads(response.choices[0].message.content)

    # ─── Phase 2A: Relationships ───────────────────────────────

    async def _run_phase2_relationships(self, user_context: str, concept_summary: str, phase1: dict) -> dict:
        valid_names = [c["name"] for c in phase1.get("concepts", [])]
        tradeoff_lines = "\n".join(
            f"- {' vs '.join(t['concepts'])}: {t['description']}"
            for t in phase1.get("tradeoffs", [])
        )

        system_prompt = f"""You are a domain expert modeling causal relationships between concepts in a quantitative simulation.

Given these concepts from the "{phase1.get('domain', 'general')}" domain:

{concept_summary}

VALID CONCEPT NAMES (use ONLY these exact names in "from" and "to" fields):
{', '.join(valid_names)}

Trade-offs identified:
{tradeoff_lines}

Produce a JSON response with this structure:

{{
    "relationships": [
        {{
            "from": "<concept_name>",
            "to": "<concept_name>",
            "direction": <1 or -1>,
            "strength": <0.0 to 1.0>,
            "nature": "<enables|constrains|trades_off|amplifies|dampens|competes>",
            "mechanism": "<2-3 sentences: WHY does this relationship exist?>",
            "responseCurve": "<linear|logarithmic|threshold|exponential>",
            "truthTable": [
                {{"fromValue": <number>, "effectOnTo": <number>, "note": "<what happens>"}},
                {{"fromValue": <number>, "effectOnTo": <number>, "note": "<what happens>"}},
                {{"fromValue": <number>, "effectOnTo": <number>, "note": "<what happens>"}}
            ],
            "confidence": "<high|medium|low>",
            "evidence": "<data source or reasoning>"
        }}
    ]
}}

RULES:
1. DIRECTION: +1 means increasing 'from' increases 'to'. -1 means increasing 'from' decreases 'to'.
2. Every trade-off listed above MUST appear as a relationship with direction: -1.
3. TRUTH TABLES must have at least 3 entries using values within each concept's validRange.
4. MECHANISMS must explain WHY, not just WHAT.
5. Include ALL significant relationships. A concept should typically relate to 2-4 others.
6. Use concept names EXACTLY as provided (snake_case)."""

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_context},
            ],
            response_format={"type": "json_object"},
            temperature=0.5,
            max_tokens=8000,
        )
        return json.loads(response.choices[0].message.content)

    # ─── Phase 2B: Attractors ──────────────────────────────────

    async def _run_phase2_attractors(self, user_context: str, concept_summary: str, phase1: dict) -> dict:
        valid_names = [c["name"] for c in phase1.get("concepts", [])]
        objectives = "; ".join(phase1.get("objectives", []))

        system_prompt = f"""You are a domain expert identifying stable strategies, regimes, and real-world reference points for a quantitative simulation.

Given these concepts from the "{phase1.get('domain', 'general')}" domain:

{concept_summary}

VALID CONCEPT NAMES (use ONLY these exact names in "position" keys):
{', '.join(valid_names)}

Objectives: {objectives}

Produce a JSON response with this structure:

{{
    "attractors": [
        {{
            "name": "<stable strategy or configuration name>",
            "description": "<why this is a natural resting point>",
            "position": {{"<concept_name>": <value_in_real_units>, ...}},
            "strength": <0.0 to 1.0>,
            "evidence": "<what real-world data supports this>"
        }}
    ],
    "referencePoints": [
        {{
            "name": "<real-world example name>",
            "description": "<what this configuration represents>",
            "position": {{"<concept_name>": <value>, ...}},
            "source": "<where this data comes from>",
            "reliability": "<high|medium|low>"
        }}
    ],
    "confidenceAssessment": {{
        "overall": "<high|medium|low>",
        "weakestLink": "<which concept or relationship is least certain>",
        "dataGaps": ["<what information would improve the model>"]
    }}
}}

RULES:
1. ATTRACTORS are the MOST IMPORTANT output. Generate 4-8+ distinct attractors.
2. Each attractor MUST have positions for ALL concepts using real-world units.
3. Attractors must be DISTINCT — genuinely different approaches.
4. REFERENCE POINTS must be real configurations that exist. Cite sources.
5. Use concept names EXACTLY as provided (snake_case)."""

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_context},
            ],
            response_format={"type": "json_object"},
            temperature=0.5,
            max_tokens=6000,
        )
        return json.loads(response.choices[0].message.content)

    # ─── Helpers ───────────────────────────────────────────────

    def _build_user_context(self, problem_statement: str, options: dict) -> str:
        user_message = ""
        doc_ctx = options.get("documentContext")

        if doc_ctx:
            structured = doc_ctx.get("structuredData", [])
            text_chunks = doc_ctx.get("textChunks", [])
            has_content = len(text_chunks) > 0 or len(structured) > 0

            if has_content:
                user_message += "IMPORTANT: The user has provided detailed supporting documents. Extract concepts, relationships, and strategies DIRECTLY from the document content.\n\n"

                if structured:
                    user_message += "═══ STRUCTURED DATA ═══\n"
                    for table in structured:
                        user_message += f"\nTable: {table.get('filename', '?')} ({table.get('rowCount', '?')} rows)\n"
                        user_message += f"Columns: {', '.join(table.get('headers', []))}\n"
                        preview = table.get("preview", [])
                        if preview:
                            user_message += f"Sample: {json.dumps(preview[:5])}\n"
                    user_message += "\n"

                if text_chunks:
                    user_message += "═══ DOCUMENT CONTENT ═══\n"
                    total = 0
                    for chunk in text_chunks:
                        if total + len(chunk) > 60000:
                            break
                        user_message += f"\n{chunk}\n"
                        total += len(chunk)

                user_message += "\n═══ USER'S QUESTION / GOAL ═══\n"

        user_message += f"Analyze this problem domain:\n\n{problem_statement}"
        return user_message

    def _build_concept_summary(self, phase1: dict) -> str:
        lines = []
        for c in phase1.get("concepts", []):
            valid_range = c.get("validRange", "unknown")
            unit = c.get("unitOfMeasure", "")
            lines.append(
                f"- {c['name']} ({c.get('displayName', c['name'])}): "
                f"{c.get('description', '')} [{c.get('nature', '?')}, "
                f"{c.get('controllability', '?')}, range: {json.dumps(valid_range)}, unit: {unit}]"
            )
        return "\n".join(lines)

    def _validate_and_normalize(self, parsed: dict, problem_statement: str) -> dict:
        """Validate and normalize the merged LLM response."""
        errors = []
        warnings = []

        concepts = parsed.get("concepts", [])
        if not concepts:
            raise ValueError("LLM response missing concepts array")
        if len(concepts) < self.min_influences:
            warnings.append(f"Only {len(concepts)} concepts (minimum {self.min_influences})")

        relationships = parsed.get("relationships", [])
        if not relationships:
            raise ValueError("LLM response missing relationships array")

        tradeoff_count = sum(1 for r in relationships if r.get("direction") == -1)
        if tradeoff_count == 0:
            warnings.append("No trade-off relationships found (direction: -1).")

        concept_names = {c["name"] for c in concepts}
        for rel in relationships:
            if not rel.get("from") and rel.get("source"):
                rel["from"] = rel["source"]
            if not rel.get("to") and rel.get("target"):
                rel["to"] = rel["target"]
            if rel.get("direction") is None and rel.get("sign") is not None:
                rel["direction"] = rel["sign"]
            if rel.get("from") not in concept_names:
                errors.append(f"Relationship references unknown concept: {rel.get('from')}")
            if rel.get("to") not in concept_names:
                errors.append(f"Relationship references unknown concept: {rel.get('to')}")

        attractors = parsed.get("attractors", [])
        if len(attractors) == 0:
            warnings.append("No attractors identified.")
        elif len(attractors) < 3:
            warnings.append(f"Only {len(attractors)} attractor(s). Works best with 3+.")

        if warnings:
            print(f"[DomainAnalyzer] Validation warnings: {'; '.join(warnings)}")

        # Normalize
        return {
            "domain": parsed.get("domain", "general"),
            "problem": {
                "statement": problem_statement,
                "restatement": parsed.get("problemRestatement", problem_statement),
                "objectives": parsed.get("objectives", []),
            },
            "concepts": [self._normalize_concept(c) for c in concepts],
            "relationships": [self._normalize_relationship(r) for r in relationships],
            "tradeoffs": parsed.get("tradeoffs", []),
            "referencePoints": parsed.get("referencePoints", []),
            "attractors": parsed.get("attractors", []),
            "confidence": parsed.get("confidenceAssessment", {
                "overall": "medium", "weakestLink": "unknown", "dataGaps": []
            }),
            "validation": {"errors": errors, "warnings": warnings, "valid": len(errors) == 0},
        }

    @staticmethod
    def _normalize_concept(c: dict) -> dict:
        return {
            "name": c["name"],
            "displayName": c.get("displayName", c["name"].replace("_", " ")),
            "description": c.get("description", ""),
            "nature": c.get("nature", "force"),
            "category": c.get("category", "general"),
            "controllability": c.get("controllability", "direct"),
            "dataDictionary": {
                "definition": c.get("description", ""),
                "unitOfMeasure": c.get("unitOfMeasure", ""),
                "precision": 2,
                "dataType": "continuous",
                "validRange": c.get("validRange", [0, 100]),
                "displayRange": c.get("displayRange") or c.get("validRange", [0, 100]),
                "displayUnit": c.get("displayUnit") or c.get("unitOfMeasure", ""),
            },
            "dynamics": {
                "inertia": c.get("dynamics", {}).get("inertia", "medium"),
                "timescale": c.get("dynamics", {}).get("timescale", "months"),
                "volatility": c.get("dynamics", {}).get("volatility", 0.1),
            },
            "historical": {
                "baseline": c.get("historicalBaseline"),
                "knownConfigurations": c.get("knownConfigurations", []),
            },
            "confidence": "medium",
        }

    @staticmethod
    def _normalize_relationship(r: dict) -> dict:
        return {
            "fromConcept": r.get("from", r.get("fromConcept", "")),
            "toConcept": r.get("to", r.get("toConcept", "")),
            "direction": r.get("direction", 1),
            "strength": max(0.0, min(1.0, r.get("strength", 0.5))),
            "nature": r.get("nature", "enables"),
            "mechanism": r.get("mechanism", ""),
            "responseCurve": {
                "type": r.get("responseCurve", "linear") if isinstance(r.get("responseCurve"), str) else r.get("responseCurve", {}).get("type", "linear"),
                "truthTable": [
                    {"fromValue": t["fromValue"], "effectOnTo": t["effectOnTo"], "note": t.get("note", "")}
                    for t in (r.get("truthTable", []) if isinstance(r.get("truthTable"), list) else r.get("responseCurve", {}).get("truthTable", []))
                ],
            },
            "confidence": r.get("confidence", "medium"),
            "evidenceSources": [r["evidence"]] if r.get("evidence") else [],
            "symmetry": {"isSymmetric": False},
        }
