"""
Forge Refine Engine
====================
Real surgical editing of an existing generated project.
No regeneration — only touches files that need to change.

Flow:
  1. Load existing project from disk (reads all .ts/.tsx/.js files)
  2. Run git diff HEAD to understand what user has manually changed
  3. Ask PLANNER to understand the refinement request in context
  4. Ask ARCHITECT which files need to change (and how)
  5. Run targeted CODEGEN only on those files (NVIDIA NIM / Groq)
  6. Run VERIFIER (tsc --noEmit) on changed files only
  7. Commit the changes via git
  8. Re-zip and optionally send to Telegram

The key invariant: files the user manually edited (tracked via git diff)
are NEVER overwritten unless the refine request explicitly targets them.
"""
from __future__ import annotations

import asyncio
import json
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from forge.core.session import GenerationSession, AppPlan
from forge.core.api import GroqClient, NvidiaClient, strip_fences
from forge.agents.codegen import CODE_SYSTEM, build_code_prompt
from forge.agents.post import run_packager

# ── Data structures ───────────────────────────────────────────────────────────

@dataclass
class RefineResult:
    success: bool
    changed_files: list[str] = field(default_factory=list)
    skipped_user_edits: list[str] = field(default_factory=list)
    error: Optional[str] = None
    zip_path: Optional[Path] = None


# ── Git helpers ───────────────────────────────────────────────────────────────

def _git_available(project_dir: Path) -> bool:
    return (project_dir / ".git").exists()


def _git_user_edited_files(project_dir: Path) -> set[str]:
    """
    Return relative paths of files the user has manually changed since
    the last FORGE commit (files dirty in the working tree vs HEAD).
    Uses real `git diff --name-only HEAD`.
    """
    if not _git_available(project_dir):
        return set()
    try:
        result = subprocess.run(
            ["git", "diff", "--name-only", "HEAD"],
            cwd=str(project_dir),
            capture_output=True, text=True, timeout=10,
        )
        return {line.strip() for line in result.stdout.splitlines() if line.strip()}
    except Exception:
        return set()


def _git_commit(project_dir: Path, message: str) -> bool:
    """Stage all changes and create a commit. Returns True on success."""
    try:
        subprocess.run(["git", "add", "-A"], cwd=str(project_dir),
                       capture_output=True, timeout=15)
        result = subprocess.run(
            ["git", "commit", "-m", message],
            cwd=str(project_dir),
            capture_output=True, text=True, timeout=15,
        )
        return result.returncode == 0
    except Exception:
        return False


# ── Refinement planner ────────────────────────────────────────────────────────

REFINE_ANALYST_SYSTEM = """You are a Principal React Native engineer doing code review and feature planning.

You are given:
1. The original app plan (JSON)
2. The current file structure of the project
3. A refinement request from the developer

Your job: produce a precise surgical edit plan.

OUTPUT: Valid JSON only. No markdown.

Schema:
{
  "summary": "One sentence: what will change",
  "files_to_modify": [
    {
      "path": "src/screens/HomeScreen.tsx",
      "reason": "Needs new search bar UI and filter logic",
      "change_type": "modify"  // "modify" | "create" | "delete"
    }
  ],
  "files_to_create": [
    {
      "path": "src/components/ui/SearchBar.tsx",
      "reason": "New reusable component required by the refine request",
      "purpose": "Search bar with debounce, clear button, keyboard handling"
    }
  ],
  "files_to_delete": [],
  "new_dependencies": [],
  "context_changes": "If AppContext needs new state fields, describe them here. Empty string if none.",
  "types_changes": "If src/types/index.ts needs new interfaces, describe them. Empty string if none.",
  "warnings": ["Warn about anything that might break existing functionality"]
}

Rules:
- Minimise blast radius — only touch what must change
- Never touch files the refine request doesn't affect
- If request implies new shared component, create it and update screens that use it
- If request touches data model, also update: types/index.ts, context/AppContext.tsx, services/storage.ts
- Be precise about path — use exact paths as they appear in the file list"""


REFINE_CODEGEN_SYSTEM = """You are a Principal React Native engineer making SURGICAL edits to an existing file.

You are given:
1. The FULL CURRENT content of the file
2. A precise description of what to change
3. Context about the app

Rules:
- Keep ALL existing code that the change doesn't affect
- Preserve the existing code style, imports structure, and naming
- Add new features cleanly, integrating with existing patterns
- Do NOT rewrite the file from scratch
- Do NOT change unrelated functionality
- Do NOT remove features that weren't asked to be removed
- Maintain TypeScript types throughout

Output ONLY the complete updated file. No markdown fences. No commentary."""


async def _plan_refinement(
    session: GenerationSession,
    refine_request: str,
    groq: GroqClient,
    all_file_paths: list[str],
) -> dict:
    """Ask Groq to analyse the request and plan which files to touch."""
    plan_json = json.dumps({
        "appName": session.plan.app_name if session.plan else "Unknown",
        "screens": [s["name"] for s in (session.plan.screens if session.plan else [])],
        "dataModels": [m["name"] for m in (session.plan.data_models if session.plan else [])],
        "navigation": session.plan.navigation if session.plan else "tab",
    })

    file_list = "\n".join(f"  {p}" for p in sorted(all_file_paths))

    prompt = (
        f"Original app plan:\n{plan_json}\n\n"
        f"Current project files:\n{file_list}\n\n"
        f"Refinement request: {refine_request}\n\n"
        "Produce the surgical edit plan JSON."
    )

    return await groq.complete_json(
        messages=[
            {"role": "system", "content": REFINE_ANALYST_SYSTEM},
            {"role": "user", "content": prompt},
        ],
        model=session.settings.plan_model,
        temperature=0.15,
    )


async def _refine_file(
    project_dir: Path,
    rel_path: str,
    reason: str,
    session: GenerationSession,
    groq: GroqClient,
    nvidia: NvidiaClient,
    sibling_context: str,
) -> Optional[str]:
    """
    Load file from disk, ask NVIDIA/Groq to make the targeted change,
    write result back. Returns new content or None on failure.
    """
    full_path = project_dir / rel_path
    if not full_path.exists():
        session.log("REFINER", f"  ⚠ File not found: {rel_path}", "warning")
        return None

    current_content = full_path.read_text(encoding="utf-8")
    plan = session.plan

    prompt = (
        f"App: {plan.app_name if plan else 'App'}\n"
        f"File to edit: {rel_path}\n"
        f"Change required: {reason}\n\n"
        f"Related files context:\n{sibling_context[:1500]}\n\n"
        f"Current file content:\n{current_content}"
    )

    code, source = await nvidia.generate_code_with_fallback(
        system_prompt=REFINE_CODEGEN_SYSTEM,
        user_prompt=prompt,
        model=session.settings.code_model,
        groq_client=groq,
        groq_model=session.settings.plan_model,
    )

    if not code or len(code.strip()) < 10:
        return None

    full_path.write_text(code, encoding="utf-8")

    # Update in-memory session.files if present
    if rel_path in session.files:
        session.files[rel_path].content = code

    session.log("REFINER", f"  ✓ {rel_path} [{source.upper()}]", "success")
    return code


async def _create_file(
    project_dir: Path,
    rel_path: str,
    purpose: str,
    session: GenerationSession,
    groq: GroqClient,
    nvidia: NvidiaClient,
) -> Optional[str]:
    """Create a brand-new file using targeted codegen."""
    plan = session.plan
    file_info = {
        "path": rel_path,
        "type": "component" if "components" in rel_path else "screen" if "screens" in rel_path else "util",
        "purpose": purpose,
        "exports": [f"default {Path(rel_path).stem}"],
        "imports": [],
        "priority": 3,
    }
    arch = session.architecture or {}
    sibling_summary = "\n".join(
        f"  {p}" for p in list(session.files.keys())[:20]
    )
    prompt = build_code_prompt(plan, arch, file_info, sibling_summary)

    code, source = await nvidia.generate_code_with_fallback(
        system_prompt=CODE_SYSTEM,
        user_prompt=prompt,
        model=session.settings.code_model,
        groq_client=groq,
        groq_model=session.settings.plan_model,
    )

    if not code:
        return None

    full_path = project_dir / rel_path
    full_path.parent.mkdir(parents=True, exist_ok=True)
    full_path.write_text(code, encoding="utf-8")
    session.add_file(rel_path, code, source=source)

    session.log("REFINER", f"  ✓ Created {rel_path} [{source.upper()}]", "success")
    return code


# ── Main refine runner ────────────────────────────────────────────────────────

async def run_refine(
    project_dir: Path,
    refine_request: str,
    session: GenerationSession,
    groq: GroqClient,
    nvidia: NvidiaClient,
) -> RefineResult:
    """
    Full surgical refinement of an existing project.
    project_dir: path to the generated project on disk
    refine_request: plain English change request
    session: the original GenerationSession (may have plan/architecture)
    """
    result = RefineResult(success=False)

    if not project_dir.exists():
        result.error = f"Project directory not found: {project_dir}"
        return result

    agent_run = session.start_agent("REFINER", f"NVIDIA NIM + Groq [{session.settings.plan_model}]")
    session.log("REFINER", f'Request: "{refine_request}"')

    # Discover user-edited files (protect them)
    user_edits = _git_user_edited_files(project_dir)
    if user_edits:
        session.log("REFINER", f"  Protecting {len(user_edits)} user-edited file(s): {', '.join(list(user_edits)[:4])}")

    # Collect all project TypeScript files
    all_files = [
        str(p.relative_to(project_dir))
        for p in project_dir.rglob("*")
        if p.suffix in (".ts", ".tsx", ".js", ".jsx", ".json")
        and "node_modules" not in str(p)
        and ".expo" not in str(p)
    ]

    # Load plan from disk if session.plan is empty
    if not session.plan:
        plan_path = project_dir / ".forge" / "plan.json"
        if plan_path.exists():
            from forge.core.session import AppPlan
            session.plan = AppPlan.from_dict(json.loads(plan_path.read_text()))

    session.log("REFINER", "Analysing refinement request…")

    # Phase 1: Plan which files to touch
    try:
        edit_plan = await _plan_refinement(session, refine_request, groq, all_files)
    except Exception as e:
        result.error = f"Refinement planning failed: {e}"
        session.fail_agent(agent_run, str(e))
        return result

    summary = edit_plan.get("summary", refine_request)
    files_to_modify = edit_plan.get("files_to_modify", [])
    files_to_create = edit_plan.get("files_to_create", [])
    files_to_delete = edit_plan.get("files_to_delete", [])
    new_deps = edit_plan.get("new_dependencies", [])
    warnings = edit_plan.get("warnings", [])

    session.log("REFINER", f"Plan: {summary}")
    session.log("REFINER", f"  Modify: {len(files_to_modify)} | Create: {len(files_to_create)} | Delete: {len(files_to_delete)}")
    for w in warnings:
        session.log("REFINER", f"  ⚠ {w}", "warning")

    # Build sibling context for code gen
    sibling_context = "\n".join(
        f"  {f['path']}: {f.get('reason', '')}" for f in files_to_modify + files_to_create
    )

    # Phase 2: Modify existing files
    for file_info in files_to_modify:
        rel_path = file_info.get("path", "")
        reason = file_info.get("reason", refine_request)

        # Never overwrite user edits unless they explicitly asked to change that file
        if rel_path in user_edits and rel_path not in refine_request:
            session.log("REFINER", f"  ⏭ Skipping {rel_path} (user-edited)", "warning")
            result.skipped_user_edits.append(rel_path)
            continue

        code = await _refine_file(
            project_dir, rel_path, reason,
            session, groq, nvidia, sibling_context
        )
        if code:
            result.changed_files.append(rel_path)

    # Phase 3: Create new files
    for file_info in files_to_create:
        rel_path = file_info.get("path", "")
        purpose = file_info.get("purpose", file_info.get("reason", ""))
        code = await _create_file(project_dir, rel_path, purpose, session, groq, nvidia)
        if code:
            result.changed_files.append(rel_path)

    # Phase 4: Delete files (with confirmation guard)
    for file_info in files_to_delete:
        rel_path = file_info.get("path", "")
        full_path = project_dir / rel_path
        if full_path.exists() and rel_path not in user_edits:
            full_path.unlink()
            result.changed_files.append(f"[deleted] {rel_path}")
            session.log("REFINER", f"  🗑 Deleted {rel_path}")

    # Phase 5: Add new npm dependencies to package.json
    if new_deps:
        pkg_path = project_dir / "package.json"
        if pkg_path.exists():
            try:
                pkg = json.loads(pkg_path.read_text())
                for dep in new_deps:
                    if "@" in dep[1:]:
                        name, version = dep.rsplit("@", 1)
                        pkg.setdefault("dependencies", {})[name] = version
                    else:
                        pkg.setdefault("dependencies", {})[dep] = "latest"
                pkg_path.write_text(json.dumps(pkg, indent=2))
                session.log("REFINER", f"  ✓ Added {len(new_deps)} dep(s) to package.json")
            except Exception as e:
                session.log("REFINER", f"  ⚠ Could not update package.json: {e}", "warning")

    # Phase 6: Git commit
    if _git_available(project_dir) and result.changed_files:
        commit_msg = f"refine: {summary[:72]}"
        committed = _git_commit(project_dir, commit_msg)
        if committed:
            session.log("REFINER", f"  ✓ git commit: {commit_msg}", "success")
        else:
            session.log("REFINER", "  ⚠ git commit failed (nothing staged?)", "warning")

    # Phase 7: Re-zip
    if session.project_dir:
        _, zip_path = await run_packager(session)
        result.zip_path = zip_path

    result.success = len(result.changed_files) > 0

    session.log_success("REFINER",
        f"✓ Refine complete — {len(result.changed_files)} file(s) changed"
        + (f", {len(result.skipped_user_edits)} user edit(s) preserved" if result.skipped_user_edits else "")
    )

    session.finish_agent(
        agent_run,
        changed=len(result.changed_files),
        protected=len(result.skipped_user_edits),
    )

    return result
