"""Schema-Driven Codegen Engine for NIKAME.

Iterates over declarative models in nikame.yaml and generates 
the full Async SQLAlchemy/Pydantic/FastAPI stack using Jinja2 templates.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from jinja2 import Environment, FileSystemLoader, StrictUndefined

from nikame.codegen.base import BaseCodegen, CodegenContext
from nikame.config.schema import DataModelConfig, FieldConfig, NikameConfig
from nikame.utils.file_writer import FileWriter
from nikame.codegen.ai_synthesizer import AISynthesizer
from nikame.codegen.context_engine import ContextEngine
from nikame.utils.logger import console


class SchemaCodegen(BaseCodegen):
    """Engine to generate code from DataModelConfig."""
    NAME = "schema"
    DESCRIPTION = "SQLAlchemy/Pydantic stack generation from data models"

    TYPE_MAP = {
        "str": {"py": "str", "sa": "String"},
        "int": {"py": "int", "sa": "Integer"},
        "float": {"py": "float", "sa": "Float"},
        "bool": {"py": "bool", "sa": "Boolean"},
        "datetime": {"py": "datetime", "sa": "DateTime"},
        "dict": {"py": "dict[str, Any]", "sa": "JSON"},
        "list": {"py": "list[Any]", "sa": "JSON"},
        "any": {"py": "Any", "sa": "JSON"},
    }

    def __init__(self, ctx: CodegenContext, config: NikameConfig) -> None:
        super().__init__(ctx, config)
        self.ai = AISynthesizer()
        self.context_engine = ContextEngine()

    @classmethod
    def should_trigger(cls, active_modules: set[str], active_features: set[str]) -> bool:
        """Trigger if either postgres or neo4j or clickhouse is present AND models are defined.
        Since we don't have the config here, we trigger if a DB is present.
        """
        db_modules = {"postgres", "neo4j", "clickhouse", "qdrant"}
        return any(m in active_modules for m in db_modules)

    async def generate(self) -> list[tuple[str, str]]:
        """Render all model-based files."""
        if not self.config.models:
            return []
            
        files = []

        # 0. Base DB files (Only once)
        db_base = """from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.ext.asyncio import AsyncAttrs

class Base(AsyncAttrs, DeclarativeBase):
    pass
"""
        files.append(("app/db/base.py", db_base))
        
        session_code = """from sqlalchemy.ext.asyncio import create_async_session_maker, create_async_engine, AsyncSession
from app.core.config import settings

engine = create_async_engine(settings.DATABASE_URL, echo=True)
async_session = create_async_session_maker(engine, expire_on_commit=False, class_=AsyncSession)

async def get_db():
    async with async_session() as session:
        yield session
"""
        files.append(("app/db/session.py", session_code))

        for name, model_cfg in self.config.models.items():
            console.print(f"[info]Synthesizing database schema and API for '{name}'...[/info]")
            
            # Context for the AI
            model_fields = json.dumps(model_cfg.model_dump().get('fields', {}), indent=2)
            biz_context = self.config.description or f"A backend service managing {name}"
            
            # 1. Models via AI
            model_prompt = self.context_engine.build_sqlalchemy_prompt(name)
            model_content = self.ai.synthesize_code(
                model_prompt, 
                f"Generate a SQLAlchemy model named {name.capitalize()} with fields:\n{model_fields}",
                context=biz_context
            )
            files.append((f"app/models/{name.lower()}.py", model_content))

            # 2. Schemas via AI
            schema_prompt = (
                "You are an expert backend engineer. Write robust Pydantic schemas (Create, Read, Update) "
                f"for the '{name}' entity. Use modern Pydantic v2 syntax. Return ONLY python code."
            )
            schema_content = self.ai.synthesize_code(
                schema_prompt,
                f"Fields required:\n{model_fields}",
                context=biz_context
            )
            files.append((f"app/schemas/{name.lower()}.py", schema_content))

            # 3. Router via AI
            router_prompt = self.context_engine.build_fastapi_router_prompt(name)
            router_content = self.ai.synthesize_code(
                router_prompt,
                f"Generate full CRUD API endpoints for {name}. Include any custom business logic implied by this context: {biz_context}",
                context=biz_context
            )
            files.append((f"app/api/v1/endpoints/{name.lower()}.py", router_content))

        # 6. Seed Data Script
        seed_content = 'print("Seed script pending AI generation")\n'
        files.append(("scripts/seed_db.py", seed_content))
        
        # 7. Requirements (Test)
        # We don't have a direct way to add to manifest here easily without changing the return type
        # or having access to manifest. But we can return a requirements file or just assume 
        # the manifest collector will pick them up if we add a wiring hook.
        
        return files

        # 5. Migration (Simple skeleton)
        # In production, we assume the user runs 'alembic revision --autogenerate'
        # but we can provide a starting point.

    def _build_context(self, name: str, cfg: DataModelConfig) -> dict[str, Any]:
        """Prepare variables for Jinja."""
        fields = {}
        primary_key = "id"
        all_model_names = list(self.config.models.keys())

        # Ensure 'id' exists if not provided
        if "id" not in cfg.fields:
            fields["id"] = self._parse_field("int", primary_key=True)

        processed_relationships = {}

        # 1. Parse fields and detect relationships
        for f_name, f_val in cfg.fields.items():
            parsed_f = self._parse_field(f_val, all_model_names=all_model_names)

            if parsed_f["is_relationship"]:
                target_model = parsed_f["target_model"]
                rel_type = parsed_f["rel_type"]
                processed_relationships[f_name] = {
                    "type": rel_type,
                    "target_model": target_model,
                    "target_table": target_model.lower(),
                    "backref": f"{name.lower()}{'s' if rel_type == 'many-to-one' else ''}",
                    "id_type": "int",
                }
            else:
                fields[f_name] = parsed_f
                if parsed_f.get("primary_key"):
                    primary_key = f_name

        # 2. Add explicit relationships from config
        for r_name, r_cfg in cfg.relationships.items():
            processed_relationships[r_name] = {
                "type": r_cfg.type,
                "target_model": r_cfg.model,
                "target_table": r_cfg.model.lower(),
                "backref": r_cfg.backref or f"{name.lower()}{'s' if r_cfg.type == 'many-to-one' else ''}",
                "id_type": "int",
            }

        return {
            "model_name": name,
            "table_name": name.lower(),
            "primary_key": primary_key,
            "fields": fields,
            "relationships": processed_relationships,
            "soft_delete": cfg.soft_delete,
        }

    def _parse_field(self, field_val: FieldConfig | str, primary_key: bool = False, all_model_names: list[str] = []) -> dict[str, Any]:
        """Parse a field definition into template-ready context."""
        if isinstance(field_val, str):
            f_type_raw = field_val
            f_cfg = FieldConfig(type=field_val)
        else:
            f_type_raw = field_val.type
            f_cfg = field_val

        res = f_cfg.model_dump()
        res["is_relationship"] = False

        # Handle Relationship shorthand: category: Category
        if f_type_raw in all_model_names:
            res["is_relationship"] = True
            res["target_model"] = f_type_raw
            res["rel_type"] = "many-to-one"

        # Handle List Relationship: items: list[Product]
        elif f_type_raw.startswith("list[") and f_type_raw[5:-1] in all_model_names:
            res["is_relationship"] = True
            res["target_model"] = f_type_raw[5:-1]
            res["rel_type"] = "one-to-many"

        # Handle Enum: type="enum" with values=[...] OR shorthand: enum[val1, val2]
        elif f_type_raw == "enum" and f_cfg.values:
            res["enum_values"] = f_cfg.values
            res["python_type"] = "str"
            res["sa_type"] = "String"
        elif f_type_raw.startswith("enum["):
            vals = [v.strip() for v in f_type_raw[5:-1].split(",")]
            res["enum_values"] = vals
            res["python_type"] = "str"
            res["sa_type"] = "String"

        # Handle List Simple: images: list[str]
        elif f_type_raw.startswith("list["):
            inner_type = f_type_raw[5:-1]
            inner_ctx = self.TYPE_MAP.get(inner_type, {"py": "str", "sa": "String"})
            res["python_type"] = f"List[{inner_ctx['py']}]"
            res["sa_type"] = "JSON"

        # Handle Standard Types
        else:
            ctx = self.TYPE_MAP.get(f_type_raw, {"py": "str", "sa": "String"})
            res["python_type"] = ctx["py"]
            res["sa_type"] = ctx["sa"]

        res["default_repr"] = self._get_default_repr(res["default"])
        return res

    def _get_default_repr(self, val: Any) -> str:
        if val is None: return "None"
        if isinstance(val, str): return f"'{val}'"
        return str(val)
