"""
Agent 3: CODE GENERATOR
=======================
Generates production-quality TypeScript/React Native code for every file.
- NVIDIA NIM CodeLlama for code files
- Groq fallback if NVIDIA fails
- Static generators for config files (no LLM needed, 100% correct)
- Context-aware: each file knows about other files
"""
from __future__ import annotations

import json
import asyncio
from pathlib import Path
from typing import Optional

from forge.core.session import GenerationSession, AppPlan
from forge.core.api import GroqClient, NvidiaClient, strip_fences

CODE_SYSTEM = """You are a Principal React Native / Expo engineer at a FAANG company.
Generate PRODUCTION-QUALITY TypeScript code. This is NOT a demo — it must WORK.

ABSOLUTE RULES:
1. ZERO placeholders — no "TODO", no "FIXME", no "implement later", no "..." 
2. ZERO hardcoded fake data — all data comes from context/storage/props
3. Real AsyncStorage operations — useEffect to load, proper saves on mutations
4. Real error handling — try/catch everywhere async, meaningful error messages
5. Real TypeScript — proper interfaces, no `any`, proper return types
6. Real navigation — useNavigation<NavigationProp>(), typed route params
7. Real loading states — isLoading boolean, ActivityIndicator
8. Real empty states — show EmptyState component when list is empty
9. Real form validation — validate before submit, show inline errors
10. Real UI — StyleSheet.create(), platform-specific shadows, proper spacing

STYLE RULES:
- Use React.FC<Props> with typed props interface
- useCallback for event handlers, useMemo for derived data
- SafeAreaView at screen root
- KeyboardAvoidingView for forms
- FlatList (not ScrollView + map) for lists
- Platform.OS checks where behaviour differs
- Haptic feedback on important actions
- Date formatting with date-fns

DO NOT output markdown fences. Output RAW code only."""


def gen_package_json(plan: AppPlan, ts: bool = True) -> str:
    deps: dict = {
        "expo": "~50.0.14",
        "expo-status-bar": "~1.11.1",
        "react": "18.2.0",
        "react-native": "0.73.6",
        "@react-navigation/native": "^6.1.17",
        "@react-navigation/stack": "^6.3.29",
        "@react-navigation/bottom-tabs": "^6.5.20",
        "@react-navigation/drawer": "^6.6.15",
        "react-native-screens": "~3.29.0",
        "react-native-safe-area-context": "4.8.2",
        "react-native-gesture-handler": "~2.14.0",
        "@react-native-async-storage/async-storage": "1.21.0",
        "@expo/vector-icons": "^14.0.2",
        "react-native-reanimated": "~3.6.2",
        "date-fns": "^3.6.0",
        "uuid": "^9.0.1",
        "expo-haptics": "~12.8.1",
        "expo-notifications": "~0.27.7",
        "expo-constants": "~15.4.5",
    }
    if plan.storage == "sqlite":
        deps["expo-sqlite"] = "~13.2.1"
    if plan.auth_required:
        deps["expo-secure-store"] = "~12.8.1"
        deps["expo-auth-session"] = "~5.5.2"

    dev_deps: dict = {
        "@babel/core": "^7.24.5",
        "@types/react": "~18.2.79",
        "@types/react-native": "^0.73.0",
        "@types/uuid": "^9.0.8",
        "typescript": "^5.3.3",
    }
    if plan.settings_include_tests if hasattr(plan, "settings_include_tests") else True:
        dev_deps["@testing-library/react-native"] = "^12.4.3"
        dev_deps["jest"] = "^29.7.0"
        dev_deps["jest-expo"] = "~50.0.4"

    data = {
        "name": plan.slug,
        "version": "1.0.0",
        "main": "node_modules/expo/AppEntry.js",
        "scripts": {
            "start": "expo start",
            "android": "expo start --android",
            "ios": "expo start --ios",
            "web": "expo start --web",
            "test": "jest --watchAll",
            "test:ci": "jest --ci",
            "lint": "eslint . --ext .ts,.tsx",
            "type-check": "tsc --noEmit",
        },
        "jest": {
            "preset": "jest-expo",
            "transformIgnorePatterns": [
                "node_modules/(?!((jest-)?react-native|@react-native(-community)?)|expo(nent)?|@expo(nent)?/.*|@expo-google-fonts/.*|react-navigation|@react-navigation/.*|@unimodules/.*|unimodules|sentry-expo|native-base|react-native-svg)"
            ]
        },
        "dependencies": deps,
        "devDependencies": dev_deps,
        "private": True,
    }
    return json.dumps(data, indent=2)


def gen_app_json(plan: AppPlan) -> str:
    data = {
        "expo": {
            "name": plan.app_name,
            "slug": plan.slug,
            "version": "1.0.0",
            "orientation": "portrait",
            "icon": "./assets/icon.png",
            "userInterfaceStyle": "automatic",
            "splash": {
                "image": "./assets/splash.png",
                "resizeMode": "contain",
                "backgroundColor": plan.background_color,
            },
            "assetBundlePatterns": ["**/*"],
            "ios": {
                "supportsTablet": True,
                "bundleIdentifier": plan.package_name,
                "infoPlist": {
                    "NSUserTrackingUsageDescription": "Used to improve app experience.",
                },
            },
            "android": {
                "adaptiveIcon": {
                    "foregroundImage": "./assets/adaptive-icon.png",
                    "backgroundColor": "#ffffff",
                },
                "package": plan.package_name,
                "permissions": ["RECEIVE_BOOT_COMPLETED", "VIBRATE"],
            },
            "web": {"favicon": "./assets/favicon.png"},
            "plugins": ["expo-notifications"],
            "extra": {
                "eas": {"projectId": "your-project-id"},
            },
        }
    }
    return json.dumps(data, indent=2)


def gen_tsconfig() -> str:
    return json.dumps({
        "extends": "expo/tsconfig.base",
        "compilerOptions": {
            "strict": True,
            "baseUrl": ".",
            "paths": {"@/*": ["src/*"]},
            "jsx": "react-native",
        },
        "include": ["**/*.ts", "**/*.tsx", ".expo/types/**/*.d.ts", "expo-env.d.ts"],
    }, indent=2)


def gen_babel_config() -> str:
    return """module.exports = function (api) {
  api.cache(true);
  return {
    presets: ['babel-preset-expo'],
    plugins: [
      'react-native-reanimated/plugin',
      [
        'module-resolver',
        {
          root: ['./'],
          alias: { '@': './src' },
        },
      ],
    ],
  };
};
"""


def gen_gitignore() -> str:
    return """# Dependencies
node_modules/
.expo/
dist/
web-build/

# Build artifacts
*.jks
*.p8
*.p12
*.key
*.mobileprovision
*.orig.*

# Env files
.env
.env.local
.env.*.local

# OS
.DS_Store
Thumbs.db

# IDE
.vscode/
.idea/
*.swp

# Logs
npm-debug.*
yarn-debug.*
yarn-error.*

# Testing
coverage/

# EAS
.easignore
"""


def gen_env_example(plan: AppPlan) -> str:
    lines = [
        "# FORGE Generated — copy to .env and fill in values",
        "",
        "# App",
        f"APP_ENV=development",
        f"APP_VERSION=1.0.0",
    ]
    for api in plan.external_apis:
        lines.append(f"# {api.upper()}_API_KEY=your_key_here")
    return "\n".join(lines) + "\n"


def gen_readme(plan: AppPlan, file_paths: list[str]) -> str:
    screens = "\n".join(f"| {s['name']} | {s['purpose']} |" for s in plan.screens)
    def _model_fields(m: dict) -> str:
        return ", ".join(field["name"] for field in m.get("fields", [])[:5])
    models = "\n".join(
        f"| {m['name']} | {_model_fields(m)} |"
        for m in plan.data_models
    )
    features = "\n".join(f"- {f}" for f in plan.features)
    tree = "\n".join(f"  {p}" for p in sorted(file_paths)[:30])

    return f"""# {plan.app_name}

> {plan.tagline}

{plan.description}

## Quick Start

```bash
# 1. Install dependencies
npm install

# 2. Start Expo dev server
npx expo start

# 3. Run on device/simulator
npx expo start --android    # Android
npx expo start --ios        # iOS (macOS only)
npx expo start --web        # Web browser
```

Scan the QR code with **Expo Go** on your device.

## Screens

| Screen | Purpose |
|--------|---------|
{screens}

## Data Models

| Model | Fields |
|-------|--------|
{models}

## Features

{features}

## Architecture

- **Framework**: React Native + Expo SDK 50
- **Language**: TypeScript (strict mode)
- **Navigation**: {plan.navigation} ({', '.join(plan.tab_screens[:3])} …)
- **State**: React Context + useReducer
- **Storage**: {plan.storage}
- **UI**: Custom design system (src/constants/theme.ts)

## Project Structure

```
{tree}
  ... ({len(file_paths)} files total)
```

## Colors

| Role | Hex |
|------|-----|
| Primary | `{plan.primary_color}` |
| Accent | `{plan.accent_color}` |
| Background | `{plan.background_color}` |
| Surface | `{plan.surface_color}` |

## Scripts

```bash
npm run type-check    # TypeScript validation
npm test             # Run test suite
npm run lint         # ESLint
```

---
*Generated by [FORGE AI](https://github.com/forge-ai/forge) v2.0*
*Powered by Groq LLaMA 3.3 + NVIDIA NIM CodeLlama*
"""


# ── Code prompts ──────────────────────────────────────────────────────────────

def build_code_prompt(
    plan: AppPlan,
    arch: dict,
    file_info: dict,
    sibling_summaries: str,
) -> str:
    path = file_info["path"]
    ftype = file_info.get("type", "")
    purpose = file_info.get("purpose", "")
    exports = file_info.get("exports", [])
    imports_needed = file_info.get("imports", [])

    models_json = json.dumps(plan.data_models, indent=2)
    screens_summary = "\n".join(
        f"  {s['name']}: {s['purpose']} | action={s.get('primaryAction','')} "
        f"| data={s.get('dataFlow','')}" for s in plan.screens
    )

    # Type-specific instructions
    type_hints = {
        "screen": f"""
SCREEN-SPECIFIC:
- Import and use theme from '@/constants/theme'
- Import useAppContext from '@/context/AppContext'  
- Use useNavigation() typed with AppStackParamList
- SafeAreaView as root container
- Real FlatList if showing a list (not ScrollView+map)
- Real empty state with EmptyState component
- Real pull-to-refresh if applicable
- Header with real title, back button if stack screen
- FAB button if hasFAB: {file_info.get('hasFAB', False)}
""",
        "component": """
COMPONENT-SPECIFIC:
- Export as React.FC<Props> with typed Props interface above component
- Pure component — no direct storage access (receive via props)
- Memoize with React.memo if expensive
- Use theme constants for all colors/spacing
""",
        "context": """
CONTEXT-SPECIFIC:
- useReducer pattern with Action union type
- AsyncStorage load on mount (useEffect)
- Auto-save on state change
- Expose: state, dispatch-wrapped action functions
- Handle: add, update, delete, clear for each model
""",
        "service": """
SERVICE-SPECIFIC:
- Pure async functions, no React hooks
- Full try/catch with typed error handling
- AsyncStorage key constants at top
- CRUD operations for all data models
- Migration support (versioned storage key)
""",
        "hook": """
HOOK-SPECIFIC:
- Custom hook returning typed tuple or object
- Handles loading/error state internally
- Memoize return value
""",
        "type": """
TYPES FILE:
- All interfaces/types for this app
- Navigation param list types
- RootState type
- All data model interfaces matching the plan exactly
- Action types for context reducer
""",
        "constant": """
CONSTANTS-SPECIFIC:
- theme.ts: Colors object, Spacing object, Typography object, Shadows object
- Use the exact brand colors from the plan
- Typography: fontSize scale, fontWeight, lineHeight
- Spacing: xs=4, sm=8, md=16, lg=24, xl=32, xxl=48
""",
        "navigation": """
NAVIGATION-SPECIFIC:
- Import all screen components
- Stack + Tab/Drawer navigator as specified
- TypedNavigation with route params
- Screen options: header style, tab bar style matching theme
- Deep linking configuration
""",
    }.get(ftype, "")

    return f"""Generate complete, production-ready TypeScript code for:

FILE: {path}
TYPE: {ftype}
PURPOSE: {purpose}
EXPORTS: {exports}
IMPORTS NEEDED: {imports_needed}

{type_hints}

APP CONTEXT:
  Name: {plan.app_name}
  Tagline: {plan.tagline}
  Primary: {plan.primary_color} | Accent: {plan.accent_color}
  Background: {plan.background_color} | Surface: {plan.surface_color}
  Navigation: {plan.navigation}
  Tab screens: {plan.tab_screens}
  Stack screens: {plan.stack_screens}
  Storage: {plan.storage} | Key: "{plan.storage_key}"
  Auth: {plan.auth_method}

ALL SCREENS:
{screens_summary}

DATA MODELS:
{models_json}

OTHER PROJECT FILES (for imports):
{sibling_summaries}

CRITICAL: Output ONLY the complete raw TypeScript code for {path}.
No markdown fences. No explanations. Start with imports. End with export.
"""


def build_test_prompt(plan: AppPlan, file_info: dict, source_code: str) -> str:
    path = file_info["path"]
    ftype = file_info.get("type", "")
    return f"""Generate Jest + @testing-library/react-native tests for:

FILE: {path} (type: {ftype})
APP: {plan.app_name}

Source code to test:
{source_code[:3000]}

Write tests covering:
- Happy path (main functionality)
- Empty state
- Error state  
- User interactions (press, input)

Use: describe/it/expect, render, fireEvent, waitFor
Mock AsyncStorage: jest.mock('@react-native-async-storage/async-storage')
Mock Navigation: jest.mock('@react-navigation/native')

Output raw TypeScript test code only. No fences."""


# ── Static file map ───────────────────────────────────────────────────────────

STATIC_GENERATORS = {
    "package.json": lambda plan, _: gen_package_json(plan),
    "app.json": lambda plan, _: gen_app_json(plan),
    "tsconfig.json": lambda _, __: gen_tsconfig(),
    "babel.config.js": lambda _, __: gen_babel_config(),
    ".gitignore": lambda _, __: gen_gitignore(),
    ".env.example": lambda plan, _: gen_env_example(plan),
}


# ── Main codegen agent ────────────────────────────────────────────────────────

async def run_codegen(
    session: GenerationSession,
    groq: GroqClient,
    nvidia: NvidiaClient,
) -> None:
    plan = session.plan
    arch = session.architecture
    files_spec = arch.get("files", [])
    total = len(files_spec)

    model_nvidia = session.settings.code_model
    model_groq = session.settings.plan_model

    agent = session.start_agent("CODEGEN", f"NVIDIA NIM → {model_nvidia}")
    session.log("CODEGEN", f"Generating {total} files…")

    # Build sibling summary for context
    sibling_summaries = "\n".join(
        f"  {f['path']}: {f.get('purpose','')}" for f in files_spec[:20]
    )

    # Generate files sequentially (parallel risks rate limits)
    for i, file_info in enumerate(files_spec):
        path = file_info["path"]
        filename = Path(path).name

        if session.on_progress:
            session.on_progress(i + 1, total, path)

        session.log("CODEGEN", f"[{i+1}/{total}] {path}")

        # Static files — no LLM needed
        if filename in STATIC_GENERATORS:
            content = STATIC_GENERATORS[filename](plan, session.settings)
            session.add_file(path, content, source="static")
            session.log("CODEGEN", f"  ✓ static ({len(content)} chars)")
            continue

        # LLM generation
        prompt = build_code_prompt(plan, arch, file_info, sibling_summaries)

        code, source = await nvidia.generate_code_with_fallback(
            system_prompt=CODE_SYSTEM,
            user_prompt=prompt,
            model=model_nvidia,
            groq_client=groq,
            groq_model=model_groq,
        )

        session.add_file(path, code, source=source)
        session.log(
            "CODEGEN",
            f"  ✓ {source.upper()} ({len(code):,} chars, {code.count(chr(10))} lines)",
            "success",
        )

    # Generate README after all files known
    readme = gen_readme(plan, list(session.files.keys()))
    session.add_file("README.md", readme, source="static")
    session.log("CODEGEN", "  ✓ README.md generated")

    session.finish_agent(
        agent,
        files=session.total_files,
        lines=session.total_lines,
        nvidia=session.nvidia_files,
        groq=session.groq_files,
    )
